import numpy as np
from sklearn import tree
import os
from IPython.display import Image
import pydotplus
import pandas as pd

# 这是决策树分类的代码，写好目标文件的地址进行处理即可。决策树图片保存在python文件所在目录。
os.environ["PATH"] += os.pathsep + 'D:/Program Files (x86)/Graphviz2.38/bin/'  # 把graphviz软件导入系统环境变量
df = pd.read_excel(r'C:\Users\liutong\Documents\Tencent Files\707921012\FileRecv\整合学生.xlsx')  # 在目标文件地址读取文件
cols = [i for i in df.columns if i in ['年龄', '性别', '专业', '入学年份']]  # 设置分类属性['年龄', '性别', '专业', '入学年份']
df1 = df.where((df.notna()), df.mean()['自杀意念'])  # 目标属性的缺失值处理
nptarget = np.array(df1['自杀意念'])  # 设置目标列为自杀意念
nptarget[nptarget < df.mean()['自杀意念']] = -1  # 归一化处理，以-1，0，1表示['低', '中', '高']
nptarget[nptarget == df.mean()['自杀意念']] = 0
nptarget[nptarget > df.mean()['自杀意念']] = 1
df2 = df.replace('男', 0).replace('女', 1)  # 0--男 1--女
# print(set(np.array(df['专业'])))
# 打印出所有专业
# {'中药学', '中医', '生物信息学', '针灸推拿', '2015级针灸推拿本科', '公共事业管理（医事法律方向）', '临床药学', '2015级中药本科',
# '康复治疗学', '护理学', '药物制剂', '临床医学(5年制二系)', '临床医学(5年制三系)', '临床医学（5年制二系）', '临床医学（5年制一系）'
# , '中药', '卫生检验与检疫', '2015级医学实验技术本科', '医学影像学', '临床医学(5+3)', '临床医学（5年制三系）', '公共事业管理',
# '卫生检疫与检疫', '食品卫生与管理', '医学实验技术', '临床医学(5年制一系)', '针灸推拿学', '护理本科', '医学影像技术', '药学',
# '医学检验技术', '英语', '预防医学', '应用统计学', '预防', '中医学', '医学信息工程', '医学英语', '临床医学(儿科医学方向)(5+3)',
# '临床医学', '2015级临床本科(专升本)', '中西医临床医学', '2015级中医本科', '食品卫生与营养学', '法医学', '英语(医学英语方向)'}
# print(len(set(np.array(df['专业']))))
# ｛1-46｝用1-46来表示这46个专业，其中的2015级为专升本等特殊情况
a = 1.0  # 设置专业序号初始值
# 遍历专业的属性用1-46来替换专业名
for n in list(set(np.array(df['专业']))):
    df2 = df2.replace(n, a)
    a += 1
npdata = np.array(df2[cols])

# 创建决策树对象，使用信息熵作为依据
clf = tree.DecisionTreeClassifier(criterion='entropy')
# fit方法分类。features为feature_names=['年龄', '性别', '专业', '入学年份']，labels为class_names=['自杀意念']
clf = clf.fit(npdata.astype('int'), nptarget.astype('int'))
# 可视化['年龄', '性别', '专业', '入学年份']=['age', 'sex', 'major', 'year'],['低', '中', '高']=['low', 'mid', 'high']
dot_data = tree.export_graphviz(clf, feature_names=['age', 'sex', 'major', 'year'], class_names=['low', 'mid', 'high'],
                                filled=True)
graph = pydotplus.graph_from_dot_data(dot_data)
Image(graph.create_png())
graph.write_png('zhenghetree.png')  # 保存图像
print(clf.score(npdata.astype('int'), nptarget.astype('int')))  # 打印出准确率得分
